import gradio as gr
import json
import os
import time
import requests
from PIL import Image


class ImgToImg:
    def __init__(self, URL, INPUT_DIR, OUTPUT_DIR):
        self.URL = URL
        self.INPUT_DIR = INPUT_DIR
        self.OUTPUT_DIR = OUTPUT_DIR

    def ui(self):
        with (gr.Blocks(analytics_enabled=False) as demo):
            workflow_name = gr.Dropdown(self.get_all_workflow_files_arr("./workflows"), type="value")
            imginput = gr.Image()
            textinput= gr.Text()
            imgoutput = gr.Image()
            btn = gr.Button("开始绘画")
            btn.click(
                fn=self.generate_image,
                inputs=[workflow_name, imginput,textinput],
                outputs=[imgoutput]
            )
        return demo

    # 开始获取请求进行编码
    def start_queue(self, prompt_workflow):
        p = {"prompt": prompt_workflow}
        data = json.dumps(p).encode('utf-8')
        requests.post(self.URL, data=data)

    # 定义获取最新图像的逻辑方法，用于之后下面函数的调用
    def get_latest_image(self, folder):
        files = os.listdir(folder)
        image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        image_files.sort(key=lambda x: os.path.getmtime(os.path.join(folder, x)))
        latest_image = os.path.join(folder, image_files[-1]) if image_files else None
        return latest_image

    # 开始生成图像，前端UI定义所需变量传递给json
    def generate_image(self, *pp):
        file = "./workflows/" + str(pp[0])
        with open(file, "r", encoding="utf-8") as file_json:
            prompt = json.load(file_json)
            # prompt["10"]["inputs"]["image"] = pp[1]
            prompt["6"]["inputs"]["text"] = f"digital artwork of a {pp[2]}"

        image = Image.fromarray(pp[1])  ## 假设 input_image 是一个 NumPy 数组，表示图像的像素数据,将 NumPy 数组转换为 PIL 中的图像对象 image
        min_side = min(image.size)  # 找到了图像的宽度和高度中的最小值，并将其存储在 min_side 变量中
        scale_factor = 512 / min_side  # 计算缩放比例，使图像最小边缩放到512像素，如果图像的最小边本身大于512，则不做任何缩放
        new_size = (round(image.size[0] * scale_factor), round(image.size[1] * scale_factor))
        resized_image = image.resize(new_size)  # 调整原始图像 resized_image 就是缩放后的新图像对象
        resized_image.save(os.path.join(self.INPUT_DIR, "man.jpg"))  # 调整图像后保存到指定的INPUT_DIR变量路径

        previous_image = self.get_latest_image(self.OUTPUT_DIR)  # 推理出的最新输出图像保存到指定的OUTPUT_DIR变量路径
        self.start_queue(prompt)
        # 这是一个循环获取指定路径的最新图像，休眠一秒钟后继续循环
        while True:
            latest_image = self.get_latest_image(self.OUTPUT_DIR)
            if latest_image != previous_image:
                return latest_image
            time.sleep(1)

    # 获取文件夹下的所有工作流的文件
    def get_all_workflow_files_arr(self, workflow_path):
        files = os.listdir(workflow_path)
        workflow_arr = [f for f in files if f.lower().endswith(('.json', '.JSON'))]
        return workflow_arr
